import config
from framework.dag import DAG
from framework.manager import CheckpointManager, RecoveryManager
from framework.operator import FileSink, MapOperator, KeyByOperator, ReduceOperator, ExactlyOnceSink, RetractOperator
from framework.watermark import WatermarkGenerator
from framework.windows import EventTimeTumblingWindow
from sources.kafka import KafkaSource, KafkaSource2


def main():
    # 初始化DAG
    dag = DAG(window_size=10, allowed_lateness=3)
    checkpoint_manager = CheckpointManager()
    recovery_manager = RecoveryManager(dag, checkpoint_manager)
    # 定义Kafka源
    kafka_source = KafkaSource(
        config.KAFKA_CONFIG['bootstrap_servers'],
        config.KAFKA_CONFIG['topic'],
        config.KAFKA_CONFIG['group_id']
    )
    
    # 定义算子节点
    source = dag.add_node('source', kafka_source)
    # Map算子
    map_op = dag.add_node('map_operator', MapOperator(lambda x: {'word': x['word'], 'count': x['count'], 'timestamp': x['timestamp'],'_event_id': x['_event_id'],'_watermark':x['_watermark']}))
    # KeyBy算子
    keyby_op = dag.add_node('keyby_operator', KeyByOperator())
    
    window_op = dag.add_node('window_operator', EventTimeTumblingWindow(30))

    reduce_op = dag.add_node('reduce_operator', ReduceOperator(RetractOperator(lambda record:record, checkpoint_manager=checkpoint_manager, node_name='reduce_operator')))
    sink = dag.add_node('sink_operator', FileSink('data.txt'))
    ExactlyOncesink = dag.add_node('ExactlyOncesink_operator', ExactlyOnceSink())
    
    # 构建DAG拓扑
    dag.add_edge('source', 'map_operator')
    dag.add_edge('map_operator', 'keyby_operator')
    dag.add_edge('keyby_operator', 'window_operator')
    dag.add_edge('window_operator', 'reduce_operator')
    dag.add_edge('reduce_operator', 'ExactlyOncesink_operator')

    # 模拟数据流执行
    for record in kafka_source:
        dag.execute(record)

if __name__ == "__main__":
    main()